import numpy

class Conductance_neuron():
    """ 
    A conductance based point neuron model.
                 dV      - Sum(i, I_i ) + I_input 
                ---- = -----------------------------
                 dt               C
    I_i are active or passive conductance based channels of the form:
                I_i = g_i * (V - E_i)
        where g_i is the conductance which may depend on any number of state
        variables, and E_i is the reversal potential which also may depend
        on any of the state variables.
    I_input could be a current injection or any number of synaptic channels.
    Finally, C is the membrane capacitance.

    Units: (unless otherwise stated)
        V --> mV       potential
        I --> nA       current
        C --> pF       capacitance
        g --> nS       conductance
    """
    def __init__(self, g_leak=1, E_leak=1, C=1, tau=None, V_init=-65):
        """
        Units:
            g_leak --> nS
            E_leak --> mV
            C      --> nF
            tau    --> ms
        """
        self.g_leak = g_leak
        self.E_leak = E_leak
        # calculate the membrane capacitance/time constant, unless it is passed
        if tau != None:
            self.tau = tau # in ms
            self.C = tau * g_leak # in ms/nS --> pF
            self.C /= 1e3 # in nF
        else:
            self.C = C # in nF
            self.tau = C / g_leak # in nF/nS --> s
            self.tau *= 1e3

        # channels and inputs used by the function self.dV_dt
        self.channels = {"I_leak":self.I_leak} # a dict of functions
        self.inputs = {}   # also a dict of functions

        # sv_index and sv_function dictionaries which hold information about
        #    the state variables and the functions used to calculate them.
        self.sv_index = {'V':0}
        self.sv_function = {'V':self.dV_dt}
        self.sv_init = {'V':V_init}
        self.sv_type = {"V":numpy.float64}
        # set up the view tupples needed for self.array_w_sv_view()
        self.setup_view_tupples()
        # set up the initial value array
        self.setup_init_array()


    def register_component(self, sv_names, sv_functions, sv_inits, sv_types=None):
        """
        Called by register_current and register_input to update dictionaries
            (should not be called by itself)
        """
        #=================== EXCEPTIONS ====================================
        # ensure that there are as many sv_functions as sv_names as sv_inits
        if len(sv_names)!=len(sv_functions) or  len(sv_inits)!=len(sv_names):
            raise RuntimeError("There must be a one-to-one correspondence "+\
                             "between sv_names and sv_functions and sv_inits" )
        # set the data types to the default of numpy.float64 unless they are passed
        if sv_types == None:
            sv_types = []
            for i in range(len(sv_names)):
                sv_types.append(numpy.float64)
        else:
            if len(sv_names)!=len(sv_types):
                raise RuntimeError("There must be a one-to-one correspondence "+\
                             "between sv_names and sv_types if you specify sv_types." )
        # ensure that the names aren't already used for sv names
        for name in sv_names:
            if name in self.sv_index.keys():
                raise RuntimeError("The name %s is already in the dictionary."\
                                   % name )
        #=============================== ====================================
        
        # update the sv_index dictionary
        max_index = len(self.sv_index.keys()) - 1
        for name in sv_names:
            self.sv_index[name] = max_index + 1
            max_index += 1
            # also add the associated function to the sv_function dictionary
            self.sv_function[name] = sv_functions[sv_names.index(name)]
            # and add the initial values associated with each sv
            self.sv_init[name] = sv_inits[sv_names.index(name)] 
            # add the data type to the list of types
            self.sv_type[name] = sv_types[sv_names.index(name)]

        # set up the view tupples needed for self.array_w_sv_view()
        self.setup_view_tupples()
        # set up the initial value array
        self.setup_init_array()
            
    def setup_init_array(self):
        """
        This method sets up the initial value array
        """
        self.sv_init_array = numpy.zeros(len(self.sv_init.keys()),dtype=numpy.float64)
        for name,index in self.sv_index.iteritems():
            self.sv_init_array[index] = self.sv_init[name]

    def setup_view_tupples(self):
        """
        This method sets up the tupples needed by the self.array_w_sv_view() method
        """
        self.view_tupples = self.sv_index.keys() # we just need a list of length = sv_index.keys()
        for name,index in self.sv_index.iteritems():
            self.view_tupples[index] = (name,self.sv_type[name])
            
        

    def register_current(self,current_name, current_funciton, sv_names=None, sv_functions=None,\
                         sv_inits=None):
        """
        Used to register a current with the neuron.  It requires names and
            functions corresponding to state variables and a function which
            given a dict of state variables and time will return the 
            value of the current in units of nA.
        """
        if sv_names != None:
            self.register_component(sv_names,sv_functions,sv_inits)

        # also update the self.channels list
        self.channels[current_name] = current_function

    def register_input(self,input_name, input_function, sv_names=None, sv_functions=None,\
                        sv_inits=None):
        """
        Used to register an input with the neuron.  It requires names and
            functions corresponding to state variables and a function which
            given a dict of state variables and time will return the 
            value of the input current in units of nA.
        """
        if sv_names != None:
            self.register_component(sv_names,sv_functions,sv_inits)
        
        # also update the self.inputs list
        self.inputs[input_name] = input_function
        
    #-----------------------------------------------------------------------
    #--                INTEGRATOR CALLED FUNCTION
    #-----------------------------------------------------------------------
    def dX_dt(self,state_variables_array,t):
        """
        This is the function that the numerical integrator calls.  It takes
            a list of state variables and returns updated state variables.
        """
        new_sv = numpy.empty_like(state_variables_array) # make the result array
        sv = state_variables_array.view(self.view_tupples) # make view
        
        # run through and update all the state variables
        for name,index in self.sv_index.iteritems():
            new_sv[index] = self.sv_function[name](sv,t)[0] # index is because array.view
   
        # translate the new state variables from a dict to a list
        return new_sv

    def dV_dt(self,sv,t):
        """
        This function returns the change in voltage/second given a state variable
            dictionary and the time.
        """
        dV = 0.0
        for key in self.channels.keys():
            dV -= self.channels[key](sv,t) # in nA
        for key in self.inputs.keys():
            dV += self.inputs[key](sv,t) # in nA
        # divide by the cell capacitance
        dV /= self.C # in ms making result in nA/nF --> V/s
        return dV
        
    def I_leak(self, sv, t):
        """ returns leak current in nA. """
        return self.g_leak * (sv['V'] - self.E_leak)/1e3 # nS * V --> nA

    def V_clamp(self, sv, t):
        """
        This is a type of input to a Conductance_neuron that will simulate holding 
            the neuron in a Voltage clamp.  The resulting neuron will likely have 
            no voltage deviations but will have an input called I_clamp.
        """
        injected_I = 0
        for key in self.channels.keys():
            injected_I += self.channels[key](sv,t) # in nA
        for key in self.inputs.keys():
            if "V_clamp" not in key:
                injected_I -= self.inputs[key](sv,t) # in nA
        return injected_I

    def set_V_clamp(self, holding_V):
        """
        set the model to voltage clamp mode with the specified holding potential.
            add the voltage clamp current input if it hasn't been added already.
        """
        # if we haven't added in the voltage clamp current... add it
        if "V_clamp" not in self.inputs.keys():
            self.register_input("V_clamp",self.V_clamp)
        # set the resting membrane potential and the initial simulation voltage to
        #   holding_V
        self.E_leak = holding_V
        self.sv_init["V"] = holding_V
        self.sv_init_array[self.sv_index["V"]] = holding_V
    
    def get_sv_init_list(self):
        """
        Return a numpy.array that are the initial values of the state variables.
        """
        return self.sv_init_array

    def calculate_input_current(self, name,  state_variables_array, t):
        """
        Return the value of the current generated by the input at the
            given time with the state variables specified.
        """
        sv = state_variables_array.view(self.view_tupples) # make view
        return -self.inputs[name](sv,t)[0]
            
    def calculate_channel_current(self, name,  state_variables_array, t):
        """
        Return the value of the current generated by the channel at the
            given time with the state variables specified.
        """
        sv = state_variables_array.view(self.view_tupples) # make view
        return self.channels[name](sv,t)[0]
